# detector_server.py
# -*- coding: utf-8 -*-
import cv2
import requests
import numpy as np
import threading
import time
from flask import Flask, Response, render_template, jsonify, request
from ultralytics import YOLO  

app = Flask(__name__)

model = YOLO('C:\\Users\\DELL\\Desktop\\yolonew\\runs\\detect\\train6\\weights\\best.pt') 

RASPBERRY_STREAM_URL = 'http://192.168.43.112:8080/?action=stream'
# 全局变量
AUTO_RUNNING = False
current_target = {'x': 0, 'y': 0, 'radius': 10}
TARGET_LOCK = threading.Lock()
SERVER_URL = 'http://192.168.43.112:5000'  # 下位机服务器地址
CENTER_X = 320
CENTER_Y = 240

def gen_detection_frames():
    stream = requests.get(RASPBERRY_STREAM_URL, stream=True)
    bytes_data = b''
    for chunk in stream.iter_content(chunk_size=1024):
        bytes_data += chunk
        a = bytes_data.find(b'\xff\xd8')  # JPEG start
        b = bytes_data.find(b'\xff\xd9')  # JPEG end
        if a != -1 and b != -1:
            jpg = bytes_data[a:b+2]
            bytes_data = bytes_data[b+2:]
            frame = cv2.imdecode(np.frombuffer(jpg, dtype=np.uint8), cv2.IMREAD_COLOR)
            if frame is None:
                continue

            results = model(frame, conf=0.8, verbose=False)
            annotated_frame = results[0].plot() 

            ret, buffer = cv2.imencode('.jpg', annotated_frame)
            frame_bytes = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

def send_move_command(direction, speed, step_duration, timeout=0.5):
    try:
        payload = {'dir': direction, 'speed': speed, 'step': step_duration}
        requests.post(f"{SERVER_URL}/move", json=payload, timeout=timeout)
        print(f"[CMD] Sent {direction} | speed={speed}, step={step_duration}")
    except Exception as e:
        print(f"[WARN] Failed to send move command: {e}")

def calculate_weed_center(stream_url):
    """
    计算置信度最高的小草的中心点像素坐标
    """
    # 加载模型
    model = YOLO('C:\\Users\\DELL\\Desktop\\yolonew\\runs\\detect\\train6\\weights\\best.pt')
    
    # 创建视频捕获对象
    cap = cv2.VideoCapture(stream_url)
    
    if not cap.isOpened():
        return None, "无法打开视频流"
    
    ret, frame = cap.read()
    if not ret:
        cap.release()
        return None, "无法获取帧"
    
    # 执行目标检测
    results = model(frame, conf=0.6)
    
    # 查找置信度最高的杂草
    highest_confidence = 0
    best_weed = None
    
    if len(results[0].boxes) > 0:
        boxes = results[0].boxes.xyxy.cpu().numpy()
        confidences = results[0].boxes.conf.cpu().numpy()
        class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
        class_names = [model.names[class_id] for class_id in class_ids]
        
        for i, (box, conf, cls_id, cls_name) in enumerate(zip(boxes, confidences, class_ids, class_names)):
            if cls_name == 'weed' and conf > highest_confidence:
                highest_confidence = conf
                x1, y1, x2, y2 = map(int, box)
                # 计算中心点坐标
                center_x = (x1 + x2) // 2
                center_y = (y1 + y2) // 2 + 25
                best_weed = {
                    'center_x': center_x,
                    'center_y': center_y,
                    'confidence': conf,
                    'bbox': (x1, y1, x2, y2)
                }
    
    cap.release()
    
    if best_weed:
        return best_weed, "找到目标杂草"
    else:
        return None, "未检测到杂草"

@app.route('/')
def index():
    return render_template('index.html') 

@app.route('/detected_feed')
def detected_feed():
    return Response(gen_detection_frames(),
                    mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/auto_detect', methods=['POST'])
def auto_detect():
    global AUTO_RUNNING
    AUTO_RUNNING = True
    CENTER_X, CENTER_Y = 320, 240  # 图像中心（假设640x480）
    LASER_X, LASER_Y = 327, 267  # 激光点 = 图像中心

    try:
        print("=" * 60)
        print("[AUTO PC] Starting 3-stage auto alignment...")
        print("=" * 60)

        # ============ 阶段1：360°旋转寻找目标 ============
        print("[STAGE 1] Rotating to find target...")
        max_rotations = 40
        for i in range(max_rotations):
            if not AUTO_RUNNING:
                break
            target, msg = calculate_weed_center(RASPBERRY_STREAM_URL)
            if target:
                print(f"[STAGE 1] Target found at ({target['center_x']}, {target['center_y']})")
                break
            send_move_command('right', speed=40, step_duration=0.15)  # 小步右转
            time.sleep(0.3)
        else:
            return jsonify({'status': 'error', 'message': 'No target found after rotation.'})

        if not target:
            return jsonify({'status': 'error', 'message': 'Target lost.'})

        # ============ 阶段2：粗调（误差 ≤ 50 像素） ============
        print("[STAGE 2] Coarse alignment (≤50px)...")
        coarse_max_iter = 15
        for attempt in range(coarse_max_iter):
            if not AUTO_RUNNING:
                break
            target, _ = calculate_weed_center(RASPBERRY_STREAM_URL)
            if not target:
                return jsonify({'status': 'error', 'message': 'Target lost in coarse stage.'})

            dx = target['center_x'] - LASER_X
            dy = target['center_y'] - LASER_Y
            # distance = max(abs(dx), abs(dy))
            print(f"[COARSE] dx={dx:.1f}, dy={dy:.1f}")

            if abs(dx) <= 30 and abs(dy) <= 30:
                print("[STAGE 2] Coarse alignment succeeded.")
                break

            # 动态速度：随尝试次数增加而减小
            speed = max(5, 15 - attempt * 2)  
            step = max(0.1, 0.3 - attempt * 0.02)
            # # 动态速度：基于尝试次数和距离的双重调整
            # # 距离权重：距离越远，速度越快（0.5-1.5倍范围）
            # distance_factor = min(1, max(0.25, distance / 50))  # 归一化到0.5-1.5范围
            
            # # 尝试次数权重：随尝试次数增加而减速
            # attempt_factor = max(0.5, 1 - attempt * 0.15)  # 从1.5线性减小到0.5
            
            # # 综合调整
            # base_speed = 10  # 基础速度值
            # base_step = 0.2  # 基础步长时间
            
            # speed = max(3, int(base_speed * distance_factor * attempt_factor))
            # step = max(0.05, base_step * distance_factor * attempt_factor)

            # 先调水平（x）
            if abs(dx) > 30:
                dir_x = 'right' if dx > 0 else 'left'
                send_move_command(dir_x, speed=speed, step_duration=step)
                time.sleep(0.3)
                continue

            # 再调垂直（y）
            if abs(dy) > 30:
                dir_y = 'backward' if dy > 0 else 'forward'
                send_move_command(dir_y, speed=speed, step_duration=step)
                time.sleep(0.3)
                continue

        # ============ 阶段3：精调（误差 ≤ 10 像素） via 舵机 ============
        print("[STAGE 3] Fine alignment with servo (≤10px)...")

        fine_max_iter = 15
        current_pan_angle = 90   # 水平舵机（左右），servo_num=2
        current_tilt_angle = 90  # 竖直舵机（上下），servo_num=1
        for attempt in range(fine_max_iter):
            if not AUTO_RUNNING:
                break
            target, _ = calculate_weed_center(RASPBERRY_STREAM_URL)
            if not target:
                break

            #dx = target['center_x'] - LASER_X
            dx = LASER_X - target['center_x'] #图像取反
            dy = LASER_Y - target['center_y']     
            if abs(dx) <= 5 and abs(dy) <= 5:
                print("[STAGE 3] Fine alignment succeeded!")
                break

            # 舵机调节：dx 正 → 目标在右 → 舵机向右转（角度增大）
            if abs(dx) > 5:
                angle_delta = np.clip(dx * 0.05, -2.5, 2.5)  
                current_pan_angle = np.clip(current_pan_angle + angle_delta, 45, 135)
                angle_to_send = int(current_pan_angle)
                try:
                    requests.post(f"{SERVER_URL}/set_servo", json={'servo_num': 2, 'angle': angle_to_send})
                    print(f"[SERVO] Set angle to {current_pan_angle:.1f}° (dx={dx:.1f})")
                except Exception as e:
                    print(f"[WARN] Servo command failed: {e}")
            # === 竖直舵机（Tilt）调节 ===
            if abs(dy) > 5:
                # dy > 0: 目标在图像上方 → 激光需上移 → tilt角度减小（假设0°为下，180°为上，90°水平）
                tilt_delta = np.clip(dy * 0.05, -2.5, 2.5)  # 竖直方向更敏感，系数略小
                current_tilt_angle = np.clip(current_tilt_angle + tilt_delta, 80, 100)  # ±10° 范围
                try:
                    requests.post(f"{SERVER_URL}/set_servo", json={'servo_num': 1, 'angle': int(current_tilt_angle)})
                except Exception as e:
                    print(f"[WARN] Tilt servo (1) failed: {e}")

            # # 动态速度：随尝试次数增加而减小
            # speed = max(2, 7 - attempt * 2)  
            # step = max(0.025, 0.1 - attempt * 0.075)
            # # 再调垂直（y）
            # if abs(dy) > 5:
            #     dir_y = 'backward' if dy > 0 else 'forward'
            #     send_move_command(dir_y, speed=speed, step_duration=step)
            #     time.sleep(0.3)
            #     continue

            time.sleep(0.4)  # 等待舵机稳定+图像更新

        # 最终开启激光
        requests.get(f"{SERVER_URL}/laser_on")
        time.sleep(5)
        requests.get(f"{SERVER_URL}/laser_off")
        time.sleep(0.2)
        requests.post(f"{SERVER_URL}/set_servo", json={'servo_num': 1, 'angle': 90})
        time.sleep(0.2)
        requests.post(f"{SERVER_URL}/set_servo", json={'servo_num': 2, 'angle': 90})
        time.sleep(0.2)
        return jsonify({'status': 'success', 'message': 'Auto alignment completed.'})

    except Exception as e:
        print(f"[AUTO PC] Exception: {e}")
        return jsonify({'status': 'error', 'message': str(e)})
    finally:
        AUTO_RUNNING = False
        print("[AUTO PC] Auto sequence ended.")

@app.route('/auto_stop_detect', methods=['POST'])
def auto_stop_detect():
    """
    停止自动模式
    """
    try:
        response = requests.get(f"{SERVER_URL}/auto_stop")
        
        if response.status_code == 200:
            return jsonify({
                'status': 'success',
                'message': '自动模式已停止'
            })
        else:
            return jsonify({
                'status': 'error',
                'message': '无法停止自动模式'
            })
            
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': f'停止自动模式失败: {str(e)}'
        })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, threaded=True)